// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <memory>
#include <mutex>
#include <random>

#include <nebula/compute/api_scalar.h>
#include <nebula/compute/kernel.h>
#include <nebula/compute/kernels/common_internal.h>
#include <nebula/compute/registry.h>
#include <nebula/util/pcg_random.h>

namespace nebula::compute::internal {

    namespace {

// We use the PCG64 single-stream ("oneseq") generator because:
// - we don't need multiple streams
// - we want deterministic output for a given seed (ruling out the unique-stream
//   PCG generators)
// - the PCG64 no-stream ("fast") generator produces identical outputs for seeds
//   which differ only by their 2 low bits (for example, 0, 1, 2, 3 all produce
//   the same output).

// Generates a random floating point number in range [0, 1).
        double generate_uniform(random::pcg64_oneseq *rng) {
            // This equation is copied from numpy. It calculates `rng() / 2^64` and
            // the return value is strictly less than 1.
            static_assert(random::pcg64_oneseq::min() == 0ULL, "");
            static_assert(random::pcg64_oneseq::max() == ~0ULL, "");
            return ((*rng)() >> 11) * (1.0 / 9007199254740992.0);
        }

        using RandomState = OptionsWrapper<RandomOptions>;

        random::pcg64_oneseq MakeSeedGenerator() {
            alkaid_vendored::pcg_extras::seed_seq_from<std::random_device> seed_source;
            random::pcg64_oneseq seed_gen(seed_source);
            return seed_gen;
        }

        turbo::Status ExecRandom(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            static random::pcg64_oneseq seed_gen = MakeSeedGenerator();
            static std::mutex seed_gen_mutex;

            random::pcg64_oneseq gen;
            const RandomOptions &options = RandomState::Get(ctx);

            if (options.initializer == RandomOptions::Seed) {
                gen.seed(options.seed);
            } else {
                std::lock_guard<std::mutex> seed_gen_lock(seed_gen_mutex);
                gen.seed(seed_gen());
            }
            double *out_values = out->array_span_mutable()->get_values<double>(1);
            for (int64_t i = 0; i < batch.length; ++i) {
                out_values[i] = generate_uniform(&gen);
            }
            return turbo::OkStatus();
        }

        const FunctionDoc random_doc{
                "Generate numbers in the range [0, 1)",
                ("Generated values are uniformly-distributed, double-precision in range [0, 1).\n"
                 "Algorithm and seed can be changed via RandomOptions."),
                {},
                "RandomOptions"};

    }  // namespace

    void register_scalar_random(FunctionRegistry *registry) {
        static auto random_options = RandomOptions::defaults();
        auto random_func = std::make_shared<ScalarFunction>(
                "random", Arity::Nullary(), random_doc, &random_options, /*is_pure=*/false);
        ScalarKernel kernel{{}, float64(), ExecRandom, RandomState::init};
        kernel.null_handling = NullHandling::OUTPUT_NOT_NULL;
        KCHECK_OK(random_func->add_kernel(kernel));
        KCHECK_OK(registry->add_function(std::move(random_func)));
    }

}  // namespace nebula::compute::internal
